from traits.api import *
from traitsui.api import *
from mayavi.core.ui.api import MayaviScene, MlabSceneModel, SceneEditor
import numpy as np
import vtk
vtk.vtkOutputWindow.SetGlobalWarningDisplay(0)

class Form(HasTraits):
    scene = Instance(MlabSceneModel, ())
    buttonClear = Button('清除视景')

    def _buttonClear_fired(self):
        mlab = self.scene.mlab
        mlab.figure(self.scene.mayavi_scene)
        mlab.clf()

    fun = Enum('dipole1', 'dipole2', 'dipole_aerial','uniform_electric_field','antenna','plane_wave','multiphysics','distance')

    def traits_init(self):
        self._fun_changed()

    def _fun_changed(self):
        mlab = self.scene.mlab
        mlab.figure(self.scene.mayavi_scene)
        mlab.clf()
        match self.fun:
            case 'dipole1':
                np.seterr(divide='ignore', invalid='ignore')
                x, y, z = np.mgrid[-10:10:51j, -10:10:51j, -10:10:51j]
                r = np.sqrt(x ** 2 + y ** 2 + z ** 2)
                s = z / r ** 3
                u = 3 * z * x / r ** 5
                v = 3 * z * y / r ** 5
                w = 3 * z * z / r ** 5 - 1 / r ** 3
                mlab.contour3d(x, y, z, s)
                mlab.quiver3d(x, y, z, u, v, w)
                mlab.flow(x, y, z, u, v, w, integration_direction='both')
                src = mlab.pipeline.vector_field(x, y, z, u, v, w)
                mlab.pipeline.vector_cut_plane(src)
                mlab.pipeline.vectors(src, mask_points=20, scale_factor=10)

            case 'dipole2':
                np.seterr(divide='ignore')

                def g(x, y, z):
                    return 1 / np.sqrt(x ** 2 + y ** 2 + z ** 2)

                def phi(x, y, z):
                    d = .01
                    return g(x, y, z + d) - g(x, y, z - d)

                def e(x, y, z, d):
                    return (phi(x + d, y, z) - phi(x, y, z)) / d, (phi(x, y + d, z) - phi(x, y, z)) / d, (
                            phi(x, y, z) - phi(x, y, z + d)) / d

                x, y, z = np.mgrid[-10:10:21j, -10:10:21j, -10:10:21j]
                phiShow = phi(x, y, z)
                u, v, w = e(x, y, z, .001)
                # r = np.sqrt(x ** 2 + y ** 2 + z ** 2)
                # s = z / r ** 3
                # u = 3 * z * x / r ** 5
                # v = 3 * z * y / r ** 5
                # w = 3 * z * z / r ** 5 - 1 / r ** 3
                mlab.contour3d(x, y, z, phiShow)
                mlab.quiver3d(x, y, z, u, v, w, scale_factor=1000)
                # self.scene.mlab.flow(x,y,z,u,v,w,integration_direction='both')
                # self.scene.mlab.pipeline.vector_cut_plane(src)
                # src = self.scene.mlab.pipeline.vector_field(x,y,z,u, v, w)
                # self.scene.mlab.pipeline.vectors(src, mask_points=20, scale_factor=10)

            case 'dipole_aerial':
                np.seterr(divide='ignore')
                [theta, phi] = np.mgrid[0:np.pi:101j, 0:2 * np.pi:201j]
                r = np.sin(theta)
                x = r * np.sin(theta) * np.cos(phi)
                y = r * np.sin(theta) * np.sin(phi)
                z = r * np.cos(theta)
                s = mlab.mesh(x, y, z)

            case 'uniform_electric_field':
                x, y, z = np.mgrid[-10:10:21j, -10:10:21j, -10:10:21j]
                u = np.ones_like(x)
                v = np.zeros_like(y)
                w = np.zeros_like(z)
                mlab.quiver3d(u, v, w)
                mlab.outline()

            case 'antenna':
                np.seterr(divide='ignore')
                [theta, phi] = np.mgrid[0 * np.pi:np.pi:101j, 0:2 * np.pi:201j]
                r = np.abs(np.sin(3 * theta))
                x = r * np.sin(theta) * np.cos(phi)
                y = r * np.sin(theta) * np.sin(phi)
                z = r * np.cos(theta)
                # mlab.mesh(x, y, z)
                x = np.sin(theta) * np.cos(phi)
                y = np.sin(theta) * np.sin(phi)
                z = np.cos(theta)
                mlab.mesh(x, y, z, scalars=r)

            case 'plane_wave':
                x, y, z = np.mgrid[-10:10:20j, -10:10:20j, -10:10:20j]
                s = x + y + 2 * z
                mlab.contour3d(s, opacity=.5)
                mlab.pipeline.image_plane_widget(mlab.pipeline.scalar_field(s), plane_orientation='x_axes', slice_index=10, )
                mlab.pipeline.image_plane_widget(mlab.pipeline.scalar_field(s), plane_orientation='y_axes', slice_index=10, )
                mlab.outline()

            case 'multiphysics':
                np.seterr(divide='ignore')
                [theta, phi] = np.mgrid[.4 * np.pi:.6 * np.pi:101j, 0:2 * np.pi:201j]
                x = np.sin(theta) * np.cos(phi)
                y = np.sin(theta) * np.sin(phi)
                z = np.cos(theta)
                r1 = np.abs(np.sin(theta)) + 1
                r2 = np.abs(np.sin(theta) * np.cos(phi)) + 2
                r3 = np.abs(np.sin(theta) * np.cos(2 * phi)) + 3
                mlab.mesh(1 * x, 1 * y, 1 * z, scalars=r1)
                mlab.mesh(2 * x, 2 * y, 2 * z, scalars=r2)
                mlab.mesh(3 * x, 3 * y, 3 * z, scalars=r2 + 1)

            case 'distance':
                np.seterr(divide='ignore')
                [theta, phi] = np.mgrid[.4 * np.pi:.6 * np.pi:101j, 0:2 * np.pi:201j]
                x = np.sin(theta) * np.cos(phi)
                y = np.sin(theta) * np.sin(phi)
                z = np.cos(theta)
                r1 = np.abs(np.sin(theta)) + 1
                r2 = np.abs(np.sin(theta) * np.cos(phi)) + 2
                r3 = np.abs(np.sin(theta) * np.cos(2 * phi)) + 3
                mlab.mesh(r1 * x, r1 * y, r1 * z)
                mlab.mesh(r2 * x, r2 * y, r2 * z)
                mlab.mesh(r3 * x, r3 * y, r3 * z)

    view = View(
        Group(
            Tabbed(
                UItem('scene', editor=SceneEditor(scene_class=MayaviScene)),
            ),
            HGroup(
                UItem('fun'),
                UItem('buttonClear'),
                spring
            ),
        ),
        height=800,
        width=1000,
        resizable=True  # We need this to resize with the parent widget
    )

if __name__ == '__main__':
    Form().configure_traits()

